# -*- coding: utf-8 -*-

import pandas as pd
from sklearn.model_selection import StratifiedKFold, train_test_split


def split_data_with_label_counts(df, col, threshold):
    return df[df[col] > threshold], df[df[col] <= threshold]


def split_data(data, grouped_cols, n_splits, data_cols, label_col):
    data = data.copy()
    data['label'] = data.groupby(grouped_cols).ngroup()
    label_counts = data['label'].value_counts()
    label_counts = label_counts.reset_index()
    label_counts.columns = ['label', 'label_counts']
    data = pd.merge(data, label_counts, on='label')

    stratified_split_data, remain_split_data = split_data_with_label_counts(data, 'label_counts', n_splits - 1)
    stratified_list = []
    if stratified_split_data.shape[0] > 0:
        stratified_data = stratified_split_data[data_cols]
        stratified_label = stratified_split_data[label_col]

        skf = StratifiedKFold(n_splits=n_splits)
        for train_idx, valid_idx in skf.split(stratified_data, stratified_label):
            stratified_train_data, stratified_valid_data = stratified_data.iloc[train_idx], stratified_data.iloc[valid_idx]
            stratified_list.append((stratified_train_data, stratified_valid_data))

    return stratified_list, remain_split_data


if __name__ == "__main__":
    base_dir = './input'
    train_dir = base_dir + '/train_preliminary'
    train_click_log = pd.read_csv(train_dir + '/click_log.csv')
    # 剔除异常用户user_id=839368
    train_click_log.drop(train_click_log[train_click_log.user_id==839368].index, inplace=True)
    for feat in train_click_log.columns:
        train_click_log[feat] = train_click_log[feat].apply(lambda x: None if x == '\\N' else x)

    data_cols = ['time', 'user_id', 'creative_id', 'click_times']
    label_col = 'label'
    grouped_cols = ['time', 'creative_id', 'click_times']
    splited_data_list = []
    n_splits = 3
    remain_split_data = train_click_log
    for i in range(3, 1, -1):
        if remain_split_data.shape[0] == 0:
            break
        stratified_list, remain_split_data = split_data(remain_split_data[data_cols], grouped_cols[:i], n_splits, data_cols, label_col)
        if len(stratified_list) > 0:
            splited_data_list.append(stratified_list)
    random_list = []
    for i in range(n_splits):
        if remain_split_data.shape[0] == 0:
            break
        random_train_data, random_valid_data, _, _ = train_test_split(remain_split_data[data_cols], remain_split_data[label_col],
                                                                      test_size=1 / n_splits, random_state=i * 1024)
        random_list.append((random_train_data, random_valid_data))
    if len(random_list) > 0:
        splited_data_list.append(random_list)

    for i in range(n_splits):
        train_data = pd.concat([t[i][0] for t in splited_data_list])
        valid_data = pd.concat([t[i][1] for t in splited_data_list])

        train_data.to_csv(base_dir + '/train_{}.csv'.format(str(i)), index=False)
        valid_data.to_csv(base_dir + '/valid_{}.csv'.format(str(i)), index=False)
